from pathlib import Path
from typing import List, Union

from llama_index.core import Document, SimpleDirectoryReader
from loguru import logger


class DataLoader:
    """数据加载器 - 支持多种数据源"""

    @staticmethod
    async def load_from_dir(file_path: Union[str, Path]) -> List[Document]:
        """从文件加载数据"""
        try:
            file_path = Path(file_path)
            if not file_path.exists() or not file_path.is_dir():
                raise FileNotFoundError(f"文件不存在: {file_path}")

            logger.info(f"开始加载文件: {file_path}")
            reader = SimpleDirectoryReader(file_path)
            documents = reader.load_data(num_workers=1, show_progress=True)

            logger.info(f"成功加载 {len(documents)} 个文档")
            return documents

        except Exception as e:
            logger.exception(f"文件加载失败: {e}")
            raise

    @staticmethod
    async def load_from_wiki_page_info(wiki_page_info_dir: Union[str, Path]) -> List[Document]:
        """从wiki_page_info目录加载数据"""
        try:
            wiki_dir = Path(wiki_page_info_dir)
            if not wiki_dir.exists() or not wiki_dir.is_dir():
                raise FileNotFoundError(f"wiki_page_info目录不存在: {wiki_dir}")

            logger.info(f"开始加载wiki_page_info目录: {wiki_dir}")
            reader = SimpleDirectoryReader(
                input_dir=str(wiki_dir),
                recursive=True,  # 递归读取子目录
                required_exts=[".txt", ".md", ".html", ".json"]  # 支持常见的wiki文件格式
            )
            documents = reader.load_data(num_workers=1, show_progress=True)

            # 为wiki文档添加特殊的元数据标识
            for doc in documents:
                doc.metadata["source_type"] = "wiki_page_info"
                doc.metadata["wiki_dir"] = str(wiki_dir)

            logger.info(f"成功从wiki_page_info目录加载 {len(documents)} 个文档")
            return documents

        except Exception as e:
            logger.exception(f"wiki_page_info目录加载失败: {e}")
            raise




